package com.jarvis.core;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.log.Log;
import cn.hutool.log.LogFactory;
import com.google.common.base.Preconditions;
import com.jarvis.core.finder.PluginFinder;
import com.jarvis.core.plugin.PluginDefine;
import java.lang.instrument.Instrumentation;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import lombok.Getter;
import lombok.Setter;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.agent.builder.AgentBuilder.Listener;
import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.scaffold.TypeValidation;
import net.bytebuddy.utility.JavaModule;
import org.apache.commons.lang3.StringUtils;

/**
 * @author xuxukang
 * @date 2019-07-04 14:20
 */
public enum Jarvis {

  /**
   * Jarvis管理器单例枚举类
   */
  INSTANCE;

  private static final Log log = LogFactory.get();

  @Setter
  @Getter
  private static Set<String> disabledPlugins = new HashSet<>(8);

  static {
    String property = System.getProperty("jarvis.disable.plugins");
    if (StringUtils.isNotBlank(property)) {
      CollUtil.addAll(disabledPlugins, StringUtils.split(property, ","));
    }
  }

  // 避免jarvis重新启动
  private AtomicBoolean initialize = new AtomicBoolean(false);
  // 用于后续重新加载已经卸载的插件（一个组件可能定义多个插件）
  private Map<String, List<AgentBuilder>> agentBuildersMap = new ConcurrentHashMap<>();
  private Map<String, List<ResettableClassFileTransformer>> resetMap = new ConcurrentHashMap<>();
  // 获取默认的Java增强对象
  private Instrumentation instrumentation = ByteBuddyAgent.install();

  public synchronized void start() {
    if (initialize.get()) {
      return;
    }

    log.info("jarvis start to initialize...");
    if (CollUtil.isNotEmpty(agentBuildersMap)) {
      for (String pluginName : agentBuildersMap.keySet()) {
        restartPlugin(pluginName);
      }
    } else {
      List<PluginDefine> pluginDefines = PluginFinder.INSTANCE.findPluginDefines(true);
      for (PluginDefine pluginDefine : pluginDefines) {
        if (disabledPlugins.contains(pluginDefine.pluginName())) {
          log.warn("【warn】{}【plugin】{}", "plugin is disabled", pluginDefine.pluginName());
          continue;
        }
        installAgent(pluginDefine);
      }
    }

    log.info("jarvis success to initialize...");
    initialize.compareAndSet(false, true);
  }

  private void installAgent(PluginDefine pluginDefine) {
    AgentBuilder agentBuilder = new AgentBuilder.Default(new ByteBuddy().with(TypeValidation.DISABLED))
        .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
        .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
        .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
        .with(new Listener() {
          @Override
          public void onDiscovery(String typeName, ClassLoader classLoader, JavaModule module,
              boolean loaded) {
            // noop
          }

        @Override
        public void onTransformation(TypeDescription typeDescription, ClassLoader classLoader,
          JavaModule module, boolean loaded, DynamicType dynamicType) {
          // noop
          log.info("【info】{}【type】{}", "onTransformation", typeDescription.getName());
        }

        @Override
        public void onIgnored(TypeDescription typeDescription, ClassLoader classLoader,
          JavaModule module, boolean loaded) {
          // noop
        }

        @Override
        public void onError(String typeName, ClassLoader classLoader, JavaModule module,
          boolean loaded, Throwable throwable) {
          log.error(throwable, "【error】{}【typeName】{}", "onError", typeName);
        }

        @Override
        public void onComplete(String typeName, ClassLoader classLoader, JavaModule module,
          boolean loaded) {
          // noop
        }
      })
      .ignore(pluginDefine.ignore())
      .type(pluginDefine.typeMatcher())
      .transform((builder, typeDescription, classLoader, module) ->
        pluginDefine.define(typeDescription, builder, classLoader)
      );
    cacheAgentBuilder(pluginDefine.pluginName(), agentBuilder);

    ResettableClassFileTransformer transformer = agentBuilder
      .installOn(instrumentation);
    cacheTransformer(pluginDefine.pluginName(), transformer);
  }

  public synchronized void stop() {
    if (!initialize.get()) {
      return;
    }

    for (String pluginName : resetMap.keySet()) {
      stopPlugin(pluginName);
    }
  }

  /**
   * 重新启动某个组件（必须在start之后使用）
   */
  public void restartPlugin(String pluginName) {
    if (!initialize.get()) {
      log.warn("【warn】{}", "Jarvis is not initialized...");
      return;
    }

    Preconditions.checkArgument(StringUtils.isNotBlank(pluginName));
    log.info("【method】{}【plugin】{}",
      "Jarvis.restartPlugin", pluginName);
    if (resetMap.containsKey(pluginName)) {// 已经启动
      log.warn("【warn】{}【pluginName】{}",
        "the plugin has been started...", pluginName);
      return;
    }

    for (AgentBuilder agentBuilder : agentBuildersMap.get(pluginName)) {
      ResettableClassFileTransformer transformer = agentBuilder
        .installOn(instrumentation);
      cacheTransformer(pluginName, transformer);
    }
  }

  public void stopPlugin(String pluginName) {
    Preconditions.checkArgument(StringUtils.isNotBlank(pluginName));
    log.info("【method】{}【plugin】{}",
      "Jarvis.stopPlugin", pluginName);
    if (CollUtil.isEmpty(resetMap.get(pluginName))) { // 没有启动
      return;
    }

    for (ResettableClassFileTransformer transformer : resetMap.get(pluginName)) {
      transformer.reset(instrumentation, AgentBuilder.RedefinitionStrategy.REDEFINITION);
    }
    resetMap.remove(pluginName);
  }

  private void cacheAgentBuilder(String pluginName, AgentBuilder agentBuilder) {
    if (agentBuildersMap.get(pluginName) == null) {
      List<AgentBuilder> agentBuilders = new ArrayList<>();
      agentBuildersMap.put(pluginName, agentBuilders);
    }
    agentBuildersMap.get(pluginName).add(agentBuilder);
  }

  private void cacheTransformer(String pluginName,
    ResettableClassFileTransformer resettableClassFileTransformer) {
    if (resetMap.get(pluginName) == null) {
      List<ResettableClassFileTransformer> resettableClassFileTransformers = new ArrayList<>();
      resetMap.put(pluginName, resettableClassFileTransformers);
    }
    resetMap.get(pluginName).add(resettableClassFileTransformer);
  }

}
